package com.winit.schedule.job.tracker.repository.mongo;

import java.util.Arrays;
import java.util.List;

import com.google.code.morphia.query.Query;
import com.google.code.morphia.query.UpdateOperations;
import com.winit.schedule.core.model.JobID;
import com.winit.schedule.core.repository.TaskItemRepository;
import com.winit.schedule.core.repository.constants.TaskItemStatus;
import com.winit.schedule.core.repository.entity.TaskItemEntity;
import com.winit.schedule.core.repository.entity.TaskRunnerEntity;
import com.winit.schedule.core.utils.Clock;
import com.winit.schedule.store.mongo.AbstractMongoRepository;

public class MongoTaskItemRepository extends AbstractMongoRepository<TaskItemEntity> implements TaskItemRepository {

    @Override
    public TaskItemEntity saveTaskItem(TaskItemEntity entity) {

        Query<TaskItemEntity> query = createQuery().field("taskItemId").equal(entity.getTaskItemId());
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .set("taskId", entity.getTaskId())
            .set("parameter", entity.getParameter())
            .set("taskItemStatus", entity.getTaskItemStatus())
            .set("lastStatusUpdateTime", entity.getLastStatusUpdateTime())
            .set("nextRunTime", entity.getNextRunTime())
            .set("autoRecovery", entity.isAutoRecovery());
        return ds.findAndModify(query, ops, false, true);
    }

    @Override
    public List<TaskItemEntity> saveTaskItem(List<TaskItemEntity> taskEntitys) {
        for (TaskItemEntity entity : taskEntitys) {
            saveTaskItem(entity);
        }
        return taskEntitys;
    }

    @Override
    public TaskItemEntity findAndFetchTaskItemByTaskId(String taskTrackerUuid, String taskId, long nextRunTime,
                                                       int fetchNum) {
        long currentTime = Clock.getSystemTime();
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId").equal(taskId);

        query.or(query.and(query.criteria("taskItemStatus").in(Arrays.asList(new TaskItemStatus[] {
                TaskItemStatus.INACTIVIE, TaskItemStatus.COMPLETE })),
            query.criteria("nextRunTime").lessThanOrEq(currentTime)),

            // 并停止或者因为异常终止的任务，下一次会被立即获取执行
            query.criteria("taskItemStatus").equal(TaskItemStatus.STOP));

        query.field("nextRunTime").greaterThan(0);

        // query.field("taskItemStatus").equal(TaskItemStatus.INACTIVIE).filter("nextRunTime <=",
        // currentTime);

        query.order("lastFetchTime,parameter").offset(0).limit(fetchNum);
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .disableValidation()
            .set("taskItemStatus", TaskItemStatus.PUSH)
            .set("nextRunTime", nextRunTime)
            .set("lastFetchTaskTrackerUuid", taskTrackerUuid)
            .set("lastFetchTime", currentTime)
            .set("lastStatusUpdateTime", currentTime);
        return ds.findAndModify(query, ops);
    }

    @Override
    public List<TaskItemEntity> getTaskItemByTaskId(String taskId) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId").equal(taskId).order("parameter");
        return query.asList();
    }

    @Override
    public void updateTaskRunnerForTaskItem(TaskRunnerEntity runner) {
        Query<TaskItemEntity> query = createQuery().field("taskItemId").in(runner.getTaskItemIds());
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .disableValidation()
            .set("taskItemStatus", TaskItemStatus.ACTIVIE)
            .set("taskRunner", runner)
            .set("lastStatusUpdateTime", Clock.getSystemTime());
        ds.update(query, ops);
    }

    @Override
    public void updateTaskRunnerForTaskItem(List<TaskRunnerEntity> runners) {
        for (TaskRunnerEntity runner : runners) {
            updateTaskRunnerForTaskItem(runner);
        }
    }

    @Override
    public void removeTaskItemByTaskId(String taskId) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId").equal(taskId);
        ds.delete(query);
    }

    @Override
    public void removeTaskRunnerByTaskTrackerUuid(String taskTrackerId) {
        Query<TaskItemEntity> query = createQuery();
        query.field("lastFetchTaskTrackerUuid")
            .equal(taskTrackerId)
            .field("taskItemStatus")
            .in(TaskItemStatus.RUN_STATUS);
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .disableValidation()
            .set("taskItemStatus", TaskItemStatus.STOP)
            // 让因连接断开导致没执行完成的任务能重新执行
            .unset("taskRunner");
        ds.update(query, ops);
    }

    /**
     * 完成任务时，将任务状态重新置为INACTIVITE，删除关联的taskRunner
     */
    @Override
    public void completeTaskItemByTaskRunnerUuid(String taskRunnerUuid, List<String> taskItemIds, TaskItemStatus status) {
        Query<TaskItemEntity> query = createQuery();
        if (null != taskItemIds) {
            query.or(query.criteria("taskItemid").in(taskItemIds),
                query.criteria("taskRunner.taskRunnerUuid").equal(taskRunnerUuid));
        } else {
            query.field("taskRunner.taskRunnerUuid").equal(taskRunnerUuid);
        }
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .disableValidation()
            .set("taskItemStatus", status)
            .unset("taskRunner");

        ds.update(query, ops);
    }

    @Override
    public void recoveryPushFailureTaskItem(long timeout) {
        Query<TaskItemEntity> query = createQuery().field("taskItemStatus")
            .in(TaskItemStatus.RUN_STATUS)
            .field("lastStatusUpdateTime")
            .lessThanOrEq(Clock.getSystemTime() - timeout);
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .disableValidation()
            .set("taskItemStatus", TaskItemStatus.STOP)
            .unset("taskRunner");
        ds.update(query, ops);
    }

    @Override
    public List<TaskItemEntity> findRunningTaskItemsByTaskId(String taskId) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId")
            .equal(taskId)
            .field("taskItemStatus")
            .in(TaskItemStatus.RUN_STATUS);
        return query.asList();
    }

    @Override
    public long findRunningTaskItemNumByTaskId(String taskId) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId")
            .equal(taskId)
            .field("taskItemStatus")
            .in(TaskItemStatus.RUN_STATUS);
        return query.countAll();
    }

    @Override
    public long findRunningTaskItemNumByJobId(List<String> taskIds) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId")
            .in(taskIds)
            .field("taskItemStatus")
            .in(TaskItemStatus.RUN_STATUS);
        return query.countAll();
    }

    @Override
    public void updateNextRunTime(JobID jobId, List<String> taskIds, long nextRunTime) {
        Query<TaskItemEntity> query = createQuery().field("taskId.taskId").in(taskIds);
        UpdateOperations<TaskItemEntity> ops = ds.createUpdateOperations(TaskItemEntity.class)
            .set("nextRunTime", nextRunTime)
            .set("taskId.jobId", jobId);
        ds.update(query, ops);
    }

    @Override
    public long queyTaskItemNumByTaskId(String taskId, long currentTime) {
        Query<TaskItemEntity> query = createQuery();
        query.field("taskId.taskId").equal(taskId);
        query.field("nextRunTime").greaterThan(0);
        query.or(query.criteria("taskItemStatus").notEqual(TaskItemStatus.COMPLETE),
            query.and(query.criteria("nextRunTime").lessThanOrEq(currentTime),
                query.criteria("taskItemStatus").equal(TaskItemStatus.COMPLETE)));
        return query.countAll();
    }

}
